package com.yifeng.repo.controller.kafka.message;

import com.gomcarter.frameworks.base.common.AssertUtils;
import com.gomcarter.frameworks.base.exception.CustomException;
import com.google.common.base.Strings;
import com.talkyun.utils.revised.Looper;
import com.yifeng.repo.base.api.TimeRange;
import com.yifeng.repo.base.utils.common.BaseUtil;
import com.yifeng.repo.controller.kafka.configure.KafkaProperties;
import com.yifeng.repo.controller.kafka.message.constant.KafkaConstant;
import com.yifeng.repo.controller.kafka.message.consume.SubscribePullConsumerLooper;
import com.yifeng.repo.controller.kafka.message.consume.SubscribePullConsumerNoPausedLooper;
import com.yifeng.repo.controller.kafka.message.handler.KafkaHelper;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.clients.consumer.internals.NoOpConsumerRebalanceListener;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.config.SaslConfigs;
import org.apache.kafka.common.security.auth.SecurityProtocol;
import org.apache.kafka.common.serialization.StringDeserializer;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import static com.yifeng.repo.controller.kafka.message.constant.KafkaConstant.AUTO_OFFSET_RESET_EARLIEST;
import static com.yifeng.repo.controller.kafka.message.constant.KafkaConstant.SASL_MECHANISM_VALUE;

/**
 * Created by daibing on 2024/4/29.
 */
@Slf4j
public class KafkaPullConsumeManager {
    protected final ConcurrentMap<String, Looper> topicSubExpression2PullConsumerLooper = new ConcurrentHashMap<>();
    protected final String applicationName;
    protected final KafkaProperties properties;

    public KafkaPullConsumeManager(String applicationName, KafkaProperties properties) {
        this.applicationName = applicationName;
        this.properties = properties;
    }

    public void init() {
        log.info("yfcloud controller Kafka PullConsumeManager init success: OWNER_ID={}, OWNER_NAME={}", BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME);
    }

    public void destroy() {
        topicSubExpression2PullConsumerLooper.forEach((key, value) -> value.shutdown());
        log.info("yfcloud controller Kafka PullConsumeManager destroy success: OWNER_ID={}, OWNER_NAME={}", BaseUtil.OWNER_ID, BaseUtil.OWNER_NAME);
    }

    /**
     * 增加Topic消费（使用订阅模式：SubscribePullConsumerLooper）
     *
     * @param topic                     消息主题
     * @param subExpression             子表达式，比如：全部是 *，多个tag是 TagA || TagC || TagD
     * @param timeRangeList             指定时间区间
     * @param consumerRebalanceListener 指定消费再平衡策略
     * @param msgListConsumer           拉取消息的消费者
     */
    public void addTopicConsumer(String topic, String subExpression, List<TimeRange> timeRangeList,
                                 ConsumerRebalanceListener consumerRebalanceListener,
                                 Consumer<ConsumerRecords<String, String>> msgListConsumer) {
        String key = getKey(topic, subExpression);
        Looper looper = topicSubExpression2PullConsumerLooper.get(key);
        if (looper != null) {
            return;
        }
        synchronized (topicSubExpression2PullConsumerLooper) {
            looper = topicSubExpression2PullConsumerLooper.get(key);
            if (looper != null) {
                return;
            }
            looper = timeRangeList == null || timeRangeList.isEmpty()
                    ? buildSubscribePullConsumerNoPausedLooper(topic, subExpression, consumerRebalanceListener, msgListConsumer)
                    : buildSubscribePullConsumerLooper(topic, subExpression, timeRangeList, consumerRebalanceListener, msgListConsumer);
            looper.start();
            topicSubExpression2PullConsumerLooper.put(key, looper);
        }
    }

    public void addTopicConsumer(String topic, String subExpression,
                                 Consumer<ConsumerRecords<String, String>> msgListConsumer) {
        // List<TimeRange> timeRangeList = TimeRange.parse("00:00~23:59");
        this.addTopicConsumer(topic, subExpression, null, new NoOpConsumerRebalanceListener(), msgListConsumer);
    }

    public void addTopicConsumer(String topic, Consumer<ConsumerRecords<String, String>> msgListConsumer) {
        String subExpression = KafkaConstant.SUB_EXPRESSION_ALL;
        // List<TimeRange> timeRangeList = TimeRange.parse("00:00~23:59");
        this.addTopicConsumer(topic, subExpression, null, new NoOpConsumerRebalanceListener(), msgListConsumer);
    }

    /**
     * 移除Topic消费
     */
    public void removeTopicConsumer(String topic, String subExpression) {
        String key = getKey(topic, subExpression);
        Looper looper = topicSubExpression2PullConsumerLooper.remove(key);
        if (looper != null) {
            looper.shutdown();
            log.info("removeTopicConsumer success: {}", looper.getName());
        }
    }

    /**
     * 列出当前消息主题的消费subExpression（订阅模式）
     */
    public List<String> listConsumerSubExpression(String topic) {
        return topicSubExpression2PullConsumerLooper.keySet().stream()
                .map(s -> s.substring((topic + "-").length()))
                .collect(Collectors.toList());
    }

    /**
     * 查询指定topic的分区信息
     */
    public List<PartitionInfo> listPartitionInfo(String topic) {
        KafkaConsumer<String, String> pullConsumer = getConsumerFromLooper(topic);
        return new ArrayList<>(pullConsumer.partitionsFor(topic));
    }

    /**
     * 获取分区的统计数据
     *
     * @param topic     消息主题
     * @param timestamp 开始时间，单位毫秒
     * @return partition -> offset
     */
    public Map<Integer, Long> listPartitionStats(String topic, long timestamp) {
        KafkaConsumer<String, String> pullConsumer = getConsumerFromLooper(topic);
        List<PartitionInfo> partitionInfos = pullConsumer.partitionsFor(topic);
        Map<TopicPartition, Long> timestampsToSearch = partitionInfos.stream().collect(Collectors.toMap(
                s -> new TopicPartition(s.topic(), s.partition()),
                s -> timestamp));
        Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes = pullConsumer.offsetsForTimes(timestampsToSearch);
        return offsetsForTimes.entrySet().stream().collect(Collectors.toMap(
                s -> s.getKey().partition(),
                s -> s.getValue().offset()
        ));
    }

    protected KafkaConsumer<String, String> getConsumerFromLooper(String topic) {
        AssertUtils.notEmpty(topicSubExpression2PullConsumerLooper.values(), "当前还没有初始化消费者，不能获取消息分区列表！");
        Looper looper = topicSubExpression2PullConsumerLooper.entrySet()
                .stream()
                .filter(entry -> entry.getKey().startsWith(topic + "-"))
                .findFirst()
                .orElseThrow(() -> new CustomException("获取不到Topic对应的消费者：" + topic)).getValue();
        return ((SubscribePullConsumerLooper) looper).getPullConsumer();
    }

    protected Looper buildSubscribePullConsumerLooper(String topic, String subExpression, List<TimeRange> timeRangeList,
                                                      ConsumerRebalanceListener consumerRebalanceListener,
                                                      Consumer<ConsumerRecords<String, String>> msgListConsumer) {
        String key = getKey(topic, subExpression);
        KafkaConsumer<String, String> pullConsumer = buildSubscribePullConsumer(topic, subExpression, consumerRebalanceListener);
        return new SubscribePullConsumerLooper(key + "-looper", 0, 1000,
                pullConsumer, timeRangeList, msgListConsumer);
    }

    protected Looper buildSubscribePullConsumerNoPausedLooper(String topic, String subExpression,
                                                              ConsumerRebalanceListener consumerRebalanceListener,
                                                              Consumer<ConsumerRecords<String, String>> msgListConsumer) {
        String key = getKey(topic, subExpression);
        KafkaConsumer<String, String> pullConsumer = buildSubscribePullConsumer(topic, subExpression, consumerRebalanceListener);
        return new SubscribePullConsumerNoPausedLooper(key + "-looper", 0, 1000, pullConsumer, msgListConsumer);
    }

    /**
     * 订阅方式：不需要考虑消息分区分配问题，支持自动重平衡
     */
    protected KafkaConsumer<String, String> buildSubscribePullConsumer(String topic, String subExpression,
                                                                       ConsumerRebalanceListener consumerRebalanceListener) {
        // 获取topic配置
        KafkaProperties.TopicProperties topicProperties = KafkaHelper.getTopicProperties(properties, topic);
        String consumerGroup = getConsumerGroup(topicProperties, subExpression);

        // build pull consumer
        Properties properties = new Properties();
        properties.put(CommonClientConfigs.SECURITY_PROTOCOL_CONFIG, SecurityProtocol.SASL_PLAINTEXT.name);
        properties.put(SaslConfigs.SASL_MECHANISM, SASL_MECHANISM_VALUE);
        properties.put(SaslConfigs.SASL_JAAS_CONFIG, KafkaHelper.getConsumerAuthContent(topicProperties));
        properties.put(ConsumerConfig.CLIENT_ID_CONFIG, applicationName + "-" + topic + "-pull-subscribe-" + subExpression + "-" + consumerGroup);
        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, topicProperties.getNamesrvAddr());
        properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        properties.put(ConsumerConfig.GROUP_ID_CONFIG, consumerGroup);
        properties.put(ConsumerConfig.AUTO_OFFSET_RESET_CONFIG, AUTO_OFFSET_RESET_EARLIEST);
        properties.put(ConsumerConfig.MAX_POLL_RECORDS_CONFIG, topicProperties.getPullBatchSize());
        properties.put(ConsumerConfig.INTERCEPTOR_CLASSES_CONFIG, topicProperties.getConsumerInterceptorClass().getName());
        properties.put(KafkaConstant.PROPERTY_TAGS, subExpression);

        // consumer start.
        KafkaConsumer<String, String> pullConsumer = new KafkaConsumer<>(properties);
        pullConsumer.subscribe(Collections.singletonList(topic), consumerRebalanceListener);
        log.info("buildSubscribePullConsumer success: {}, consumerGroup={}, pullBatchSize={}",
                getKey(topic, subExpression), consumerGroup, topicProperties.getPullBatchSize());
        return pullConsumer;
    }

    /**
     * 指定订阅表达式使用指定消费组，如果没有指定订阅表达式，或者如果没有配置，都使用默认消费组
     */
    protected String getConsumerGroup(KafkaProperties.TopicProperties topicProperties, String subExpression) {
        return Strings.isNullOrEmpty(subExpression) || topicProperties.getConsumerGroupMapping() == null
                ? topicProperties.getConsumerGroup()
                : topicProperties.getConsumerGroupMapping().getOrDefault(subExpression, topicProperties.getConsumerGroup());
    }

    private String getKey(String topic, String subExpression) {
        return topic + "-" + subExpression;
    }

}
